import torch.nn as nn
import torch


class CNNForEEG(nn.Module):
    def __init__(self):
        super(CNNForEEG, self).__init__()
        self.conv1 = nn.Conv2d(
            in_channels=10,
            out_channels=32,
            kernel_size=5,
            padding=2,
        )
        self.conv2 = nn.Conv2d(
            in_channels=32,
            out_channels=64,
            kernel_size=5,
            padding=2,
        )
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(
            in_channels=64,
            out_channels=128,
            kernel_size=5,
            padding=2,
        )
        self.conv4 = nn.Conv2d(
            in_channels=128,
            out_channels=256,
            kernel_size=5,
            padding=2,
        )
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()
        self.dense = nn.Linear(
            in_features=256 * 8 * 8,
            out_features=256
        )
        self.fc1 = nn.Linear(256, 40)
        self.fc2 = nn.Linear(40, 2)
        self.relu = nn.ReLU()

# 输

    def forward(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2(x1)
        x3 = self.pool1(x2)
        x4 = self.conv3(x3)
        x5 = self.conv4(x4)
        x6 = self.pool2(x5)
        x7 = self.dense(self.flatten(x6))
        x8 = self.fc1(x7)
        return self.fc2(self.relu(x8))


if __name__ == "__main__":
    model = CNNForEEG()
    # 验证输入输出尺寸
    # 模拟输入：(batch_size, channels, height, width)
    input_tensor = torch.randn(1, 10, 32, 32)
    output_tensor = model(input_tensor)
    print("输入维度:", input_tensor.shape)      # 输出: torch.Size([1, 10, 32, 32])
    print("输出维度:", output_tensor)
